package cn.com.registry.etcd;

import com.google.common.util.concurrent.ListenableFuture;
import com.google.protobuf.ByteString;
import com.ibm.etcd.api.DeleteRangeResponse;
import com.ibm.etcd.api.KeyValue;
import com.ibm.etcd.api.LeaseGrantResponse;
import com.ibm.etcd.api.RangeResponse;
import com.ibm.etcd.client.KvStoreClient;
import com.ibm.etcd.client.kv.KvClient;
import com.ibm.etcd.client.lease.LeaseClient;
import com.ibm.etcd.client.lease.PersistentLease;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * @author jiaming
 */
public class EtcdClient {
    private static final Logger log = LogManager.getLogger();

    private KvClient kvClient;
    private LeaseClient leaseClient;

    public EtcdClient(KvStoreClient kvStoreClient) {
        this.kvClient = kvStoreClient.getKvClient();
        this.leaseClient = kvStoreClient.getLeaseClient();
    }

    public boolean checkPathExists(String key) {
        return get(key) != null;
    }

    public void put(String key, String value) {
        kvClient.put(ByteString.copyFromUtf8(key), ByteString.copyFromUtf8(value)).prevKv().sync();
    }

    public String get(String key) {
        RangeResponse rangeResponse = kvClient.get(ByteString.copyFromUtf8(key)).sync();
        List<KeyValue> keyValues = rangeResponse.getKvsList();
        if (keyValues == null || keyValues.isEmpty()) {
            return null;
        }
        return keyValues.get(0).getValue().toStringUtf8();
    }

    public List<KeyValue> getPrefix(String key) {
        RangeResponse rangeResponse = kvClient.get(ByteString.copyFromUtf8(key)).asPrefix().sync();
        List<KeyValue> keyValues = rangeResponse.getKvsList();
        if (keyValues == null || keyValues.isEmpty()) {
            return null;
        }
        return keyValues;
    }

    public void delete(String key) {
        ListenableFuture<DeleteRangeResponse> result = kvClient.delete(ByteString.copyFromUtf8(key)).backoffRetry().async();
    }

    public void deletePrefix(String keyPrefix) {
        ListenableFuture<DeleteRangeResponse> result = kvClient.delete(ByteString.copyFromUtf8(keyPrefix)).asPrefix().backoffRetry().async();
    }

    public KvClient.WatchIterator watch(String key) {
        return kvClient.watch(ByteString.copyFromUtf8(key)).start();
    }

    public KvClient.WatchIterator watchPrefix(String key) {
        return kvClient.watch(ByteString.copyFromUtf8(key)).asPrefix().start();
    }

    public long putAndGrant(String key, String value, long ttl) {
        LeaseGrantResponse lease = leaseClient.grant(ttl).sync();
        put(key, value, lease.getID());
        return lease.getID();
    }

    public void put(String key, String value, long id) {
        kvClient.put(ByteString.copyFromUtf8(key), ByteString.copyFromUtf8(value), id).sync();
    }

    public long keepAlive(String key, String value, int frequencySecs, int minTtl) throws Exception {
        PersistentLease lease = leaseClient.maintain()
                .keepAliveFreq(frequencySecs)
                .minTtl(minTtl)
                .start();
        long newId = lease.get(5L, TimeUnit.SECONDS);
        put(key, value, newId);
        return newId;
    }

}